from langchain_community.embeddings import HuggingFaceBgeEmbeddings
import torch

#Model Download
from modelscope import snapshot_download
model_dir = snapshot_download('BAAI/bge-large-zh-v1.5')

model_kwargs = {"device":"cuda"}
encode_kwargs = {"normalize_embeddings":True}

hf = HuggingFaceBgeEmbeddings(model_name = model_dir,
                              model_kwargs = model_kwargs,
                              encode_kwargs = encode_kwargs,
                              query_instruction = "为这个句子生成表示以用于检索相关文章："
                              )

queries = ["最新的AI研究成果","健康饮食的重要性"]
passages = ["AI技术正在不断进步，最新的研究揭示了其在医疗领域的潜在应用。", "合理的饮食习惯对维持良好的身体健康至关重要，包括足够的蔬菜和水果。"]

q_embeddings = torch.stack([torch.tensor(hf.embed_query(query)) for query in queries])
print(q_embeddings.shape)

p_embeddings = torch.tensor(hf.embed_documents(passages))
print(p_embeddings.shape)


scores = q_embeddings @ p_embeddings.T
print(scores)